#
#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

import logging
import copy
import time
import os

from huggingface_hub import snapshot_download

from api.utils.file_utils import get_project_base_directory
from rag.settings import PARALLEL_DEVICES
from .operators import *  # noqa: F403
from . import operators
import math
import numpy as np
import cv2
import onnxruntime as ort

from .postprocess import build_post_process

loaded_models = {}

# 定义一个 transform 函数，用于依次应用一系列操作（ops）到数据上
def transform(data, ops=None):
    """ transform """
    # 如果 ops 为 None，则初始化为空列表
    if ops is None:
        ops = []
    # 遍历每个操作 op，并将 data 传入执行
    for op in ops:
        data = op(data)
        # 如果某个操作返回 None，则整个 transform 返回 None
        if data is None:
            return None
    # 返回经过所有操作处理后的数据
    return data


# 定义 create_operators 函数，根据参数列表创建操作符对象
def create_operators(op_param_list, global_config=None):
    """
    create operators based on the config

    Args:
        params(list): a dict list, used to create some operators
    """
    # 确保 op_param_list 是一个列表类型
    assert isinstance(
        op_param_list, list), ('operator config should be a list')
    # 初始化空的操作符列表
    ops = []
    # 遍历每个 operator 配置项
    for operator in op_param_list:
        # 确保每个 operator 是字典类型且只包含一个键值对
        assert isinstance(operator,
                          dict) and len(operator) == 1, "yaml format error"
        # 提取操作符名称（即字典的 key）
        op_name = list(operator)[0]
        # 获取对应的参数，若无参数则设为空字典
        param = {} if operator[op_name] is None else operator[op_name]
        # 如果提供了全局配置，则合并进参数中
        if global_config is not None:
            param.update(global_config)
        # 使用 getattr 从 operators 模块获取对应名称的操作符类并实例化
        op = getattr(operators, op_name)(**param)
        # 将实例化的操作符加入 ops 列表
        ops.append(op)
    # 返回构建好的操作符列表
    return ops


# 定义 load_model 函数，用于加载 ONNX 模型文件并缓存
def load_model(model_dir, nm, device_id: int | None = None):
    # 构造模型文件路径
    model_file_path = os.path.join(model_dir, nm + ".onnx")
    # 根据设备 ID 创建缓存标签以区分不同设备上的模型
    model_cached_tag = model_file_path + str(device_id) if device_id is not None else model_file_path

    # 声明使用全局变量 loaded_models 来缓存已加载的模型
    global loaded_models
    # 尝试从缓存中获取模型
    loaded_model = loaded_models.get(model_cached_tag)
    # 如果缓存命中，则记录日志并返回已有模型
    if loaded_model:
        logging.info(f"load_model {model_file_path} reuses cached model")
        return loaded_model

    # 检查模型文件是否存在，不存在则抛出异常
    if not os.path.exists(model_file_path):
        raise ValueError("not find model file path {}".format(
            model_file_path))

    # 内部函数判断 CUDA 是否可用
    def cuda_is_available():
        try:
            import torch
            if torch.cuda.is_available() and torch.cuda.device_count() > device_id:
                return True
        except Exception:
            return False
        return False

    # 创建 ONNX Runtime 的 SessionOptions 实例
    options = ort.SessionOptions()
    # 关闭内存池以减少内存占用
    options.enable_cpu_mem_arena = False
    # 设置顺序执行模式
    options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
    # 设置线程数
    options.intra_op_num_threads = 2
    options.inter_op_num_threads = 2

    # https://github.com/microsoft/onnxruntime/issues/9509#issuecomment-951546580
    # Shrink GPU memory after execution
    # 执行后收缩 GPU 内存
    run_options = ort.RunOptions()
    # 判断是否使用 CUDA
    if cuda_is_available():
        # CUDA 提供商选项设置
        # 使用特定 GPU
        # 限制 GPU 内存
        # GPU 内存分配策略
        cuda_provider_options = {
            "device_id": device_id, # Use specific GPU
            "gpu_mem_limit": 512 * 1024 * 1024, # Limit gpu memory
            "arena_extend_strategy": "kNextPowerOfTwo",  # gpu memory allocation strategy
        }
        # 创建使用 CUDA 的推理会话
        sess = ort.InferenceSession(
            model_file_path,
            options=options,
            providers=['CUDAExecutionProvider'],
            provider_options=[cuda_provider_options]
            )
        # 添加运行时配置，启用 GPU 内存收缩
        run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", "gpu:" + str(device_id))
        logging.info(f"load_model {model_file_path} uses GPU")
    else:
        # 创建使用 CPU 的推理会话
        sess = ort.InferenceSession(
            model_file_path,
            options=options,
            providers=['CPUExecutionProvider'])
        # 添加运行时配置，启用 CPU 内存收缩
        run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", "cpu")
        logging.info(f"load_model {model_file_path} uses CPU")
    # 构建模型和运行选项的元组
    loaded_model = (sess, run_options)
    # 缓存模型以便后续复用
    loaded_models[model_cached_tag] = loaded_model
    # 返回加载的模型
    return loaded_model


# 定义 TextRecognizer 类，用于文本识别任务，支持图像预处理、模型推理和后处理流程
class TextRecognizer:
    def __init__(self, model_dir, device_id: int | None = None):
        # 初始化识别图像的输入形状，默认为 [3, 48, 320]（通道数、高度、宽度）
        self.rec_image_shape = [int(v) for v in "3, 48, 320".split(",")]
        # 设置识别阶段的批量大小为16
        self.rec_batch_num = 16
        # 配置后处理参数，使用 CTC 解码器，并指定字符字典路径和是否使用空格字符
        postprocess_params = {
            'name': 'CTCLabelDecode',
            "character_dict_path": os.path.join(model_dir, "ocr.res"),
            "use_space_char": True
        }
        # 构建后处理操作对象
        self.postprocess_op = build_post_process(postprocess_params)
        # 加载 ONNX 模型文件和运行配置
        self.predictor, self.run_options = load_model(model_dir, 'rec', device_id)
        # 获取模型输入张量的描述信息
        self.input_tensor = self.predictor.get_inputs()[0]

    # 图像缩放并归一化处理函数，根据最大宽高比调整图像尺寸
    def resize_norm_img(self, img, max_wh_ratio):
        # 获取图像输入形状参数：通道数、高度、宽度
        imgC, imgH, imgW = self.rec_image_shape

        # 断言输入图像的通道数与设定一致
        assert imgC == img.shape[2]
        # 根据最大宽高比重新计算图像宽度
        imgW = int((imgH * max_wh_ratio))
        # 获取模型输入张量的第三个维度（宽度），若为字符串或非正整数则忽略
        w = self.input_tensor.shape[3:][0]
        if isinstance(w, str):
            pass
        elif w is not None and w > 0:
            imgW = w
        # 获取原始图像的高度和宽度
        h, w = img.shape[:2]
        # 计算宽高比
        ratio = w / float(h)
        # 如果按比例缩放后的宽度超过限制，则使用最大宽度；否则使用计算出的宽度
        if math.ceil(imgH * ratio) > imgW:
            resized_w = imgW
        else:
            resized_w = int(math.ceil(imgH * ratio))

        # 对图像进行缩放
        resized_image = cv2.resize(img, (resized_w, imgH))
        # 转换数据类型为 float32 并将通道顺序转为 CHW
        resized_image = resized_image.astype('float32')
        resized_image = resized_image.transpose((2, 0, 1)) / 255
        # 归一化处理：减去均值 0.5，再除以标准差 0.5
        resized_image -= 0.5
        resized_image /= 0.5
        # 创建一个全零的填充图像张量
        padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
        # 将缩放后的图像放入填充图像中
        padding_im[:, :, 0:resized_w] = resized_image
        # 返回填充后的图像
        return padding_im

    # 特定于视觉语言任务的图像缩放与归一化处理函数
    def resize_norm_img_vl(self, img, image_shape):

        # 解析输入图像形状参数
        imgC, imgH, imgW = image_shape
        # BGR 转 RGB
        img = img[:, :, ::-1]  # bgr2rgb
        # 缩放图像到指定大小
        resized_image = cv2.resize(
            img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
        # 数据类型转换和通道顺序转换
        resized_image = resized_image.astype('float32')
        resized_image = resized_image.transpose((2, 0, 1)) / 255
        # 返回归一化后的图像
        return resized_image

    # SRN 模型专用的图像缩放与归一化函数
    def resize_norm_img_srn(self, img, image_shape):
        # 解析图像形状参数
        imgC, imgH, imgW = image_shape

        # 创建一个黑色背景图像
        img_black = np.zeros((imgH, imgW))
        # 获取输入图像的高度和宽度
        im_hei = img.shape[0]
        im_wid = img.shape[1]

        # 根据图像宽高比选择不同的缩放策略
        if im_wid <= im_hei * 1:
            img_new = cv2.resize(img, (imgH * 1, imgH))
        elif im_wid <= im_hei * 2:
            img_new = cv2.resize(img, (imgH * 2, imgH))
        elif im_wid <= im_hei * 3:
            img_new = cv2.resize(img, (imgH * 3, imgH))
        else:
            img_new = cv2.resize(img, (imgW, imgH))

        # 将图像转换为 numpy 数组
        img_np = np.asarray(img_new)
        # 转换为灰度图
        img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
        # 将灰度图像复制到黑色背景上
        img_black[:, 0:img_np.shape[1]] = img_np
        # 增加一个通道维度
        img_black = img_black[:, :, np.newaxis]

        # 获取图像尺寸并设置通道数为1
        row, col, c = img_black.shape
        c = 1

        # 重塑图像形状为 (channel, height, width)，并转换为 float32
        return np.reshape(img_black, (c, row, col)).astype(np.float32)

    # 定义 srn_other_inputs 方法，用于生成 SRN 模型所需的额外输入数据
    def srn_other_inputs(self, image_shape, num_heads, max_text_length):

        # 解析图像形状参数：通道数、高度、宽度
        imgC, imgH, imgW = image_shape
        # 计算特征维度大小，基于图像高度和宽度除以8后的乘积
        feature_dim = int((imgH / 8) * (imgW / 8))

        # 创建编码器位置信息数组，范围从0到feature_dim-1，并转换为int64类型
        encoder_word_pos = np.array(range(0, feature_dim)).reshape(
            (feature_dim, 1)).astype('int64')
        # 创建 GS（Global Sequence）模块的位置信息数组，范围从0到max_text_length-1，并转换为int64类型
        gsrm_word_pos = np.array(range(0, max_text_length)).reshape(
            (max_text_length, 1)).astype('int64')

        # 初始化注意力偏差矩阵，所有元素初始化为1.0
        gsrm_attn_bias_data = np.ones((1, max_text_length, max_text_length))
        # 计算上三角注意力偏差矩阵，并调整形状后重复num_heads次，乘以负无穷大(-1e9)，确保在softmax时不被选中
        gsrm_slf_attn_bias1 = np.triu(gsrm_attn_bias_data, 1).reshape(
            [-1, 1, max_text_length, max_text_length])
        gsrm_slf_attn_bias1 = np.tile(
            gsrm_slf_attn_bias1,
            [1, num_heads, 1, 1]).astype('float32') * [-1e9]

        # 类似地计算下三角注意力偏差矩阵
        gsrm_slf_attn_bias2 = np.tril(gsrm_attn_bias_data, -1).reshape(
            [-1, 1, max_text_length, max_text_length])
        gsrm_slf_attn_bias2 = np.tile(
            gsrm_slf_attn_bias2,
            [1, num_heads, 1, 1]).astype('float32') * [-1e9]

        # 将位置信息数组增加一个维度，以便于后续处理
        encoder_word_pos = encoder_word_pos[np.newaxis, :]
        gsrm_word_pos = gsrm_word_pos[np.newaxis, :]

        # 返回编码器位置信息、GS位置信息、以及上下三角注意力偏差矩阵
        return [
            encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1,
            gsrm_slf_attn_bias2
        ]

    # 定义 process_image_srn 方法，用于处理输入图像并准备SRN模型需要的全部输入数据
    def process_image_srn(self, img, image_shape, num_heads, max_text_length):
        # 调用 resize_norm_img_srn 方法对图像进行缩放和归一化处理，并添加一个新的维度
        norm_img = self.resize_norm_img_srn(img, image_shape)
        norm_img = norm_img[np.newaxis, :]

        # 调用 srn_other_inputs 方法获取其他必要的输入数据
        [encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1, gsrm_slf_attn_bias2] = \
            self.srn_other_inputs(image_shape, num_heads, max_text_length)

        # 确保返回的数据类型正确
        gsrm_slf_attn_bias1 = gsrm_slf_attn_bias1.astype(np.float32)
        gsrm_slf_attn_bias2 = gsrm_slf_attn_bias2.astype(np.float32)
        encoder_word_pos = encoder_word_pos.astype(np.int64)
        gsrm_word_pos = gsrm_word_pos.astype(np.int64)

        # 返回所有处理好的输入数据
        return (norm_img, encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1,
                gsrm_slf_attn_bias2)

    # SAR 模型专用图像缩放与归一化函数，支持宽度自适应调整并填充
    def resize_norm_img_sar(self, img, image_shape,
                            width_downsample_ratio=0.25):
        # 解析输入图像形状配置：通道数、高度、最小宽度、最大宽度
        imgC, imgH, imgW_min, imgW_max = image_shape
        # 获取当前图像的高度 h 和宽度 w
        h = img.shape[0]
        # 获取当前图像的宽度 w
        w = img.shape[1]
        # 初始化有效比例为 1.0，表示未被裁剪或压缩
        valid_ratio = 1.0
        # make sure new_width is an integral multiple of width_divisor.
        # 计算宽度除数，确保缩放后宽度是下采样率的整数倍
        width_divisor = int(1 / width_downsample_ratio)
        # resize
        # 根据图像宽高比计算目标宽度
        ratio = w / float(h)
        # 计算目标宽度 resize_w，保持高度不变，按比例缩放宽度
        resize_w = math.ceil(imgH * ratio)
        # 确保目标宽度是 width_divisor 的整数倍，避免无法分割的问题
        if resize_w % width_divisor != 0:
            resize_w = round(resize_w / width_divisor) * width_divisor
        # 如果设置了最小宽度，则取最大值以保证不低于限制
        if imgW_min is not None:
            resize_w = max(imgW_min, resize_w)
        # 如果设置了最大宽度，则限制宽度不超过最大值，并更新有效比例
        if imgW_max is not None:
            valid_ratio = min(1.0, 1.0 * resize_w / imgW_max)
            resize_w = min(imgW_max, resize_w)
        # 使用 OpenCV 将图像缩放到目标尺寸 (resize_w, imgH)
        resized_image = cv2.resize(img, (resize_w, imgH))
        # 转换为 float32 类型以便后续归一化处理
        resized_image = resized_image.astype('float32')
        # norm
        # 如果图像是灰度图（通道数为1），则进行单通道归一化处理
        if image_shape[0] == 1:
            # 归一化到 [0,1] 区间，并增加一个维度作为通道维度
            resized_image = resized_image / 255
            resized_image = resized_image[np.newaxis, :]
        else:
            # 对于彩色图像，将通道顺序从 HWC 转换为 CHW，并归一化到 [0,1]
            resized_image = resized_image.transpose((2, 0, 1)) / 255
        # 对图像进行标准化处理：减去均值 0.5，再除以标准差 0.5，映射到 [-1,1]
        resized_image -= 0.5
        resized_image /= 0.5
        # 获取经过缩放后的图像形状
        resize_shape = resized_image.shape
        # 创建一个填充图像张量，使用 -1 填充未使用的区域
        padding_im = -1.0 * np.ones((imgC, imgH, imgW_max), dtype=np.float32)
        # 将实际缩放后的图像放入填充张量中，其余部分保留填充值
        padding_im[:, :, 0:resize_w] = resized_image
        # 获取填充后的图像形状
        pad_shape = padding_im.shape

        # 返回填充后的图像、缩放后的形状、填充后的形状以及有效比例
        return padding_im, resize_shape, pad_shape, valid_ratio

    # SPIN 模型专用图像缩放与归一化函数，适用于单通道灰度图像
    def resize_norm_img_spin(self, img):
        # 将图像转换为灰度图
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        # return padding_im
        # 固定缩放图像至 100x32 大小，采用立方插值法
        img = cv2.resize(img, tuple([100, 32]), cv2.INTER_CUBIC)
        # 转换为 float32 类型
        img = np.array(img, np.float32)
        # 增加一个维度，使其变为 (1, 32, 100)
        img = np.expand_dims(img, -1)
        # 将图像维度从 HWC 转换为 CHW，即 (1, 32, 100)
        img = img.transpose((2, 0, 1))
        # 设置均值为 [127.5]，用于标准化
        mean = [127.5]
        # 设置标准差为 [127.5]，用于标准化
        std = [127.5]
        # 转换为 numpy 数组并调整形状为 (1, 1)
        mean = np.array(mean, dtype=np.float32)
        std = np.array(std, dtype=np.float32)
        mean = np.float32(mean.reshape(1, -1))
        stdinv = 1 / np.float32(std.reshape(1, -1))
        # 减去均值并乘以标准差倒数，实现标准化到 [-1,1]
        img -= mean
        img *= stdinv
        # 返回标准化后的图像
        return img

    # SVTR 模型专用图像缩放与归一化函数，适用于通用图像处理
    def resize_norm_img_svtr(self, img, image_shape):

        # 解析图像形状参数：通道数、高度、宽度
        imgC, imgH, imgW = image_shape
        # 使用线性插值方法将图像缩放至指定大小
        resized_image = cv2.resize(
            img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
        # 转换为 float32 类型
        resized_image = resized_image.astype('float32')
        # 将图像维度从 HWC 转换为 CHW，并归一化到 [0,1]
        resized_image = resized_image.transpose((2, 0, 1)) / 255
        # 对图像进行标准化处理：减去均值 0.5，再除以标准差 0.5，映射到 [-1,1]
        resized_image -= 0.5
        resized_image /= 0.5
        # 返回标准化后的图像
        return resized_image

    # ABINet 模型专用的图像缩放和归一化方法
    def resize_norm_img_abinet(self, img, image_shape):

        # 解析输入图像的形状配置：通道数、高度、宽度
        imgC, imgH, imgW = image_shape

        # 使用线性插值法将图像缩放到指定大小
        resized_image = cv2.resize(
            img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
        resized_image = resized_image.astype('float32')
        # 将像素值从 [0, 255] 归一化到 [0, 1]
        resized_image = resized_image / 255.

        # 定义均值和标准差，用于标准化处理
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        # 标准化处理：减去均值并除以标准差
        resized_image = (
            resized_image - mean[None, None, ...]) / std[None, None, ...]
        resized_image = resized_image.transpose((2, 0, 1))
        # 再次确保数据类型为 float32
        resized_image = resized_image.astype('float32')

        return resized_image

    # CAN 模型专用的图像归一化方法
    def norm_img_can(self, img, image_shape):

        # 将图像转换为灰度图，因为 CAN 只能处理灰度图像
        img = cv2.cvtColor(
            img, cv2.COLOR_BGR2GRAY)  # CAN only predict gray scale image

        if self.rec_image_shape[0] == 1:
            h, w = img.shape
            _, imgH, imgW = self.rec_image_shape
            # 若原始图像小于目标尺寸，则进行零填充
            if h < imgH or w < imgW:
                padding_h = max(imgH - h, 0)
                padding_w = max(imgW - w, 0)
                img_padded = np.pad(img, ((0, padding_h), (0, padding_w)),
                                    'constant',
                                    constant_values=(255))
                img = img_padded

        # 将图像扩展维度，并归一化至 [0, 1]
        img = np.expand_dims(img, 0) / 255.0  # h,w,c -> c,h,w
        # 类型转换
        img = img.astype('float32')

        return img

    # 主函数，用于批量处理图像列表
    def __call__(self, img_list):
        # 获取图像列表中的图像数量
        img_num = len(img_list)
        # Calculate the aspect ratio of all text bars
        # 计算所有文本条目的宽高比
        width_list = []
        for img in img_list:
            width_list.append(img.shape[1] / float(img.shape[0]))
        # Sorting can speed up the recognition process
        # 对宽高比进行排序以加快识别过程
        indices = np.argsort(np.array(width_list))
        rec_res = [['', 0.0]] * img_num
        batch_num = self.rec_batch_num
        st = time.time()

        # 批量处理图像
        for beg_img_no in range(0, img_num, batch_num):
            end_img_no = min(img_num, beg_img_no + batch_num)
            norm_img_batch = []
            imgC, imgH, imgW = self.rec_image_shape[:3]
            max_wh_ratio = imgW / imgH
            # max_wh_ratio = 0
            for ino in range(beg_img_no, end_img_no):
                h, w = img_list[indices[ino]].shape[0:2]
                wh_ratio = w * 1.0 / h
                max_wh_ratio = max(max_wh_ratio, wh_ratio)
            for ino in range(beg_img_no, end_img_no):
                norm_img = self.resize_norm_img(img_list[indices[ino]],
                                                max_wh_ratio)
                norm_img = norm_img[np.newaxis, :]
                norm_img_batch.append(norm_img)
            norm_img_batch = np.concatenate(norm_img_batch)
            norm_img_batch = norm_img_batch.copy()

            input_dict = {}
            input_dict[self.input_tensor.name] = norm_img_batch
            for i in range(100000):
                try:
                    outputs = self.predictor.run(None, input_dict, self.run_options)
                    break
                except Exception as e:
                    if i >= 3:
                        raise e
                    time.sleep(5)
            preds = outputs[0]
            rec_result = self.postprocess_op(preds)
            for rno in range(len(rec_result)):
                rec_res[indices[beg_img_no + rno]] = rec_result[rno]

        return rec_res, time.time() - st


# 初始化TextDetector类时设置预处理操作列表
# 预处理包括调整图像大小、归一化、转换为CHW格式以及保留特定键值
# 设置模型目录和可选设备ID，并初始化预测器与运行选项
class TextDetector:
    def __init__(self, model_dir, device_id: int | None = None):
        # 定义预处理步骤，包括缩放、标准化等
        pre_process_list = [{
            'DetResizeForTest': {
                'limit_side_len': 960,
                'limit_type': "max",
            }
        }, {
            'NormalizeImage': {
                'std': [0.229, 0.224, 0.225],
                'mean': [0.485, 0.456, 0.406],
                'scale': '1./255.',
                'order': 'hwc'
            }
        }, {
            'ToCHWImage': None
        }, {
            'KeepKeys': {
                'keep_keys': ['image', 'shape']
            }
        }]
        # 定义后处理参数，用于文本检测结果的优化
        postprocess_params = {"name": "DBPostProcess", "thresh": 0.3, "box_thresh": 0.5, "max_candidates": 1000,
                              "unclip_ratio": 1.5, "use_dilation": False, "score_mode": "fast", "box_type": "quad"}

        self.postprocess_op = build_post_process(postprocess_params)
        self.predictor, self.run_options = load_model(model_dir, 'det', device_id)
        self.input_tensor = self.predictor.get_inputs()[0]

        img_h, img_w = self.input_tensor.shape[2:]
        if isinstance(img_h, str) or isinstance(img_w, str):
            pass
        elif img_h is not None and img_w is not None and img_h > 0 and img_w > 0:
            # 如果指定了输入张量的高度和宽度，则更新预处理步骤中的缩放参数
            pre_process_list[0] = {
                'DetResizeForTest': {
                    'image_shape': [img_h, img_w]
                }
            }
        self.preprocess_op = create_operators(pre_process_list)

    # 按顺时针方向排序点，确保四边形的顶点顺序一致
    def order_points_clockwise(self, pts):
        rect = np.zeros((4, 2), dtype="float32")
        s = pts.sum(axis=1)
        rect[0] = pts[np.argmin(s)]
        rect[2] = pts[np.argmax(s)]
        tmp = np.delete(pts, (np.argmin(s), np.argmax(s)), axis=0)
        diff = np.diff(np.array(tmp), axis=1)
        rect[1] = tmp[np.argmin(diff)]
        rect[3] = tmp[np.argmax(diff)]
        return rect

    # 裁剪检测结果，确保所有点都在图像边界内
    def clip_det_res(self, points, img_height, img_width):
        for pno in range(points.shape[0]):
            points[pno, 0] = int(min(max(points[pno, 0], 0), img_width - 1))
            points[pno, 1] = int(min(max(points[pno, 1], 0), img_height - 1))
        return points

    # 过滤掉过小的检测框，避免误检
    def filter_tag_det_res(self, dt_boxes, image_shape):
        img_height, img_width = image_shape[0:2]
        dt_boxes_new = []
        for box in dt_boxes:
            if isinstance(box, list):
                box = np.array(box)
            box = self.order_points_clockwise(box)
            box = self.clip_det_res(box, img_height, img_width)
            rect_width = int(np.linalg.norm(box[0] - box[1]))
            rect_height = int(np.linalg.norm(box[0] - box[3]))
            if rect_width <= 3 or rect_height <= 3:
                continue
            dt_boxes_new.append(box)
        dt_boxes = np.array(dt_boxes_new)
        return dt_boxes

    def filter_tag_det_res_only_clip(self, dt_boxes, image_shape):
        img_height, img_width = image_shape[0:2]
        dt_boxes_new = []
        for box in dt_boxes:
            if isinstance(box, list):
                box = np.array(box)
            box = self.clip_det_res(box, img_height, img_width)
            dt_boxes_new.append(box)
        dt_boxes = np.array(dt_boxes_new)
        return dt_boxes

    def __call__(self, img):
        ori_im = img.copy()
        data = {'image': img}

        st = time.time()
        data = transform(data, self.preprocess_op)
        img, shape_list = data
        if img is None:
            return None, 0
        img = np.expand_dims(img, axis=0)
        shape_list = np.expand_dims(shape_list, axis=0)
        img = img.copy()
        input_dict = {}
        input_dict[self.input_tensor.name] = img
        for i in range(100000):
            try:
                outputs = self.predictor.run(None, input_dict, self.run_options)
                break
            except Exception as e:
                if i >= 3:
                    raise e
                time.sleep(5)

        post_result = self.postprocess_op({"maps": outputs[0]}, shape_list)
        dt_boxes = post_result[0]['points']
        dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im.shape)

        return dt_boxes, time.time() - st


class OCR:
    def __init__(self, model_dir=None):
        """
        If you have trouble downloading HuggingFace models, -_^ this might help!!

        For Linux:
        export HF_ENDPOINT=https://hf-mirror.com

        For Windows:
        Good luck
        ^_-

        """
        if not model_dir:
            try:
                model_dir = os.path.join(
                        get_project_base_directory(),
                        "rag/res/deepdoc")
                
                # Append muti-gpus task to the list
                if PARALLEL_DEVICES is not None and PARALLEL_DEVICES > 0:
                    self.text_detector = []
                    self.text_recognizer = []
                    for device_id in range(PARALLEL_DEVICES):
                        self.text_detector.append(TextDetector(model_dir, device_id))
                        self.text_recognizer.append(TextRecognizer(model_dir, device_id))
                else:
                    self.text_detector = [TextDetector(model_dir, 0)]
                    self.text_recognizer = [TextRecognizer(model_dir, 0)]

            except Exception:
                model_dir = snapshot_download(repo_id="InfiniFlow/deepdoc",
                                              local_dir=os.path.join(get_project_base_directory(), "rag/res/deepdoc"),
                                              local_dir_use_symlinks=False)
                
                if PARALLEL_DEVICES is not None:
                    assert PARALLEL_DEVICES > 0, "Number of devices must be >= 1"
                    self.text_detector = []
                    self.text_recognizer = []
                    for device_id in range(PARALLEL_DEVICES):
                        self.text_detector.append(TextDetector(model_dir, device_id))
                        self.text_recognizer.append(TextRecognizer(model_dir, device_id))
                else:
                    self.text_detector = [TextDetector(model_dir, 0)]
                    self.text_recognizer = [TextRecognizer(model_dir, 0)]

        self.drop_score = 0.5
        self.crop_image_res_index = 0

    def get_rotate_crop_image(self, img, points):
        # 该函数用于对图像中任意四边形区域进行透视变换并裁剪成矩形图像，常用于 OCR 中的文本框矫正
        '''
        img_height, img_width = img.shape[0:2]
        left = int(np.min(points[:, 0]))
        right = int(np.max(points[:, 0]))
        top = int(np.min(points[:, 1]))
        bottom = int(np.max(points[:, 1]))
        img_crop = img[top:bottom, left:right, :].copy()
        points[:, 0] = points[:, 0] - left
        points[:, 1] = points[:, 1] - top
        '''
        # 确保输入的点是一个四边形（4个点，每个点2个坐标）
        assert len(points) == 4, "shape of points must be 4*2"
        # 计算目标图像的宽度：取两组对边长度的最大值
        img_crop_width = int(
            max(
                np.linalg.norm(points[0] - points[1]),
                np.linalg.norm(points[2] - points[3])))
        # 计算目标图像的高度：取另外两组对边长度的最大值
        img_crop_height = int(
            max(
                np.linalg.norm(points[0] - points[3]),
                np.linalg.norm(points[1] - points[2])))
        # 定义目标图像的标准四个顶点（左上、右上、右下、左下）
        pts_std = np.float32([[0, 0], [img_crop_width, 0],
                              [img_crop_width, img_crop_height],
                              [0, img_crop_height]])
        # 使用 cv2.getPerspectiveTransform 获取透视变换矩阵 M
        M = cv2.getPerspectiveTransform(points, pts_std)
        # 应用透视变换，将原图中的四边形区域映射为标准矩形图像
        dst_img = cv2.warpPerspective(
            img,
            M, (img_crop_width, img_crop_height),
            borderMode=cv2.BORDER_REPLICATE,
            flags=cv2.INTER_CUBIC)
        # 获取变换后的图像尺寸，用于判断是否需要旋转
        dst_img_height, dst_img_width = dst_img.shape[0:2]
        # 如果高宽比大于等于 1.5，则将图像顺时针旋转90度（便于后续识别模型处理）
        if dst_img_height * 1.0 / dst_img_width >= 1.5:
            dst_img = np.rot90(dst_img)
        # 返回最终的矫正后图像
        return dst_img

    def sorted_boxes(self, dt_boxes):
        """
        Sort text boxes in order from top to bottom, left to right
        args:
            dt_boxes(array):detected text boxes with shape [4, 2]
        return:
            sorted boxes(array) with shape [4, 2]
        """
        num_boxes = dt_boxes.shape[0]
        sorted_boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
        _boxes = list(sorted_boxes)

        for i in range(num_boxes - 1):
            for j in range(i, -1, -1):
                if abs(_boxes[j + 1][0][1] - _boxes[j][0][1]) < 10 and \
                        (_boxes[j + 1][0][0] < _boxes[j][0][0]):
                    tmp = _boxes[j]
                    _boxes[j] = _boxes[j + 1]
                    _boxes[j + 1] = tmp
                else:
                    break
        return _boxes

    def detect(self, img, device_id: int | None = None):
        if device_id is None:
            device_id = 0

        time_dict = {'det': 0, 'rec': 0, 'cls': 0, 'all': 0}

        if img is None:
            return None, None, time_dict

        start = time.time()
        dt_boxes, elapse = self.text_detector[device_id](img)
        time_dict['det'] = elapse

        if dt_boxes is None:
            end = time.time()
            time_dict['all'] = end - start
            return None, None, time_dict

        return zip(self.sorted_boxes(dt_boxes), [
                   ("", 0) for _ in range(len(dt_boxes))])

    def recognize(self, ori_im, box, device_id: int | None = None):
        if device_id is None:
            device_id = 0

        img_crop = self.get_rotate_crop_image(ori_im, box)

        rec_res, elapse = self.text_recognizer[device_id]([img_crop])
        text, score = rec_res[0]
        if score < self.drop_score:
            return ""
        return text

    # 批量识别图像中的文字内容，支持指定设备（如 GPU）进行推理
    def recognize_batch(self, img_list, device_id: int | None = None):
        # 如果未指定设备 ID，默认使用第 0 号设备（例如默认 GPU）
        if device_id is None:
            device_id = 0
        # 调用文本识别器对图像列表进行 OCR 推理，返回识别结果和耗时
        rec_res, elapse = self.text_recognizer[device_id](img_list)
        # 初始化识别结果文本列表
        texts = []
        # 遍历每一张图像的识别结果
        for i in range(len(rec_res)):
            # 获取当前图像的文字和置信度分数
            text, score = rec_res[i]
            # 如果置信度低于阈值，则将识别结果清空（视为不可靠）
            if score < self.drop_score:
                text = ""
            # 将处理后的文本加入结果列表
            texts.append(text)
        # 返回批量识别后的文本结果列表
        return texts

    def __call__(self, img, device_id = 0, cls=True):
        time_dict = {'det': 0, 'rec': 0, 'cls': 0, 'all': 0}
        if device_id is None:
            device_id = 0

        if img is None:
            return None, None, time_dict

        start = time.time()
        ori_im = img.copy()
        dt_boxes, elapse = self.text_detector[device_id](img)
        time_dict['det'] = elapse

        if dt_boxes is None:
            end = time.time()
            time_dict['all'] = end - start
            return None, None, time_dict

        img_crop_list = []

        dt_boxes = self.sorted_boxes(dt_boxes)

        for bno in range(len(dt_boxes)):
            tmp_box = copy.deepcopy(dt_boxes[bno])
            img_crop = self.get_rotate_crop_image(ori_im, tmp_box)
            img_crop_list.append(img_crop)

        rec_res, elapse = self.text_recognizer[device_id](img_crop_list)

        time_dict['rec'] = elapse

        filter_boxes, filter_rec_res = [], []
        for box, rec_result in zip(dt_boxes, rec_res):
            text, score = rec_result
            if score >= self.drop_score:
                filter_boxes.append(box)
                filter_rec_res.append(rec_result)
        end = time.time()
        time_dict['all'] = end - start

        # for bno in range(len(img_crop_list)):
        #    print(f"{bno}, {rec_res[bno]}")

        return list(zip([a.tolist() for a in filter_boxes], filter_rec_res))
